package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"sync"
	"time"
)

type DiskConf struct {
	DiskNo    int           `json:"no"`
	Disk      string        `json:"disk"`
	N         int64         `json:"n"`
	BsStr     string        `json:"bs_str"`
	Bs        int64         `json:"-"`
	ReadDelay time.Duration `json:"read_delay"`
}

var (
	disk = &DiskConf{}
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	flag.IntVar(&disk.DiskNo, "disk_no", -1, "-disk_no 3 (/dev/nvme3n1)")
	flag.StringVar(&disk.Disk, "disk", "", "-disk /dev/nvme3n1")
	flag.Int64Var(&disk.N, "n", 1, "-n 20")
	flag.StringVar(&disk.BsStr, "bs", "128K", "-bs 128K")
	flag.DurationVar(&disk.ReadDelay, "delay", 0, "-delay 1ms")
}

func main() {
	flag.Parse()
	parseConfig()
	startReadTask(disk)
	return
}

func parseConfig() {
	var err error
	if disk.Disk == "" {
		if disk.DiskNo >= 0 {
			disk.Disk = fmt.Sprintf("/dev/nvme%dn1", disk.DiskNo)
		}
	}
	if disk.Disk == "" {
		log.Fatalf("miss disk")
	}
	if disk.N == 0 {
		log.Fatalf("parallel > 0")
	}
	if disk.BsStr == "" {
		log.Fatalf("bs == 0")
	}
	disk.Bs, err = parseBytes(disk.BsStr)
	if err != nil {
		log.Fatalf("bs not invalid. err=%v", err)
	}
	if disk.ReadDelay < 0 {
		log.Fatalf("delay must > 0")
	}
}

var (
	bytesM = map[uint8]int64{
		'B': 1,
		'K': 1 << 10,
		'M': 1 << 20,
		'G': 1 << 30,
		'T': 1 << 40,
	}
)

func parseBytes(str string) (int64, error) {
	if str == "" {
		return 0, errors.New("read_bytes is empty")
	}

	base, ok := bytesM[str[len(str)-1]]
	if !ok {
		return 0, errors.New("read_bytes need B, K, M, G, T")
	}
	num, err := strconv.Atoi(str[:len(str)-1])
	if err != nil {
		return 0, err
	}
	return base * int64(num), nil
}

func startRead(disk *DiskConf, f *os.File, wg *sync.WaitGroup, buffer []byte, off int64) {
	_, err := f.ReadAt(buffer, off)
	if err != nil && err != io.EOF {
		log.Fatalf("read[%s] err=%v", disk.Disk, err)
	}
	wg.Done()
}

//go:linkname nanotime1 runtime.nanotime1
func nanotime1() int64

func startReadTask(disk *DiskConf) {
	f, err := os.OpenFile(disk.Disk, os.O_RDONLY, os.ModeDevice)
	if err != nil {
		log.Fatalf("disk[%s] open failed. err=%v", disk.Disk, err)
	}

	var (
		wg      sync.WaitGroup
		buffers = make([][]byte, disk.N)
		off     = int64(0)
	)

	for i := int64(0); i < disk.N; i++ {
		buffers[i] = make([]byte, disk.Bs)
	}

	devSize, err := f.Seek(0, io.SeekEnd)
	if err != nil {
		log.Fatalf("disk[%s] get size. err=%v", disk.Disk, err)
	}

	log.Printf("start read task disk[%s] size[%d]", disk.Disk, devSize)

	_, err = f.Seek(0, io.SeekStart)
	if err != nil {
		log.Fatalf("disk[%s] seek start. err=%v", disk.Disk, err)
	}

	for off < devSize {
		begin := nanotime1()
		for i := int64(0); i < disk.N; i++ {
			wg.Add(1)
			go startRead(disk, f, &wg, buffers[i], off)
			off += disk.Bs
		}
		wg.Wait()
		end := nanotime1()

		if time.Duration(end-begin)*time.Nanosecond < disk.ReadDelay {
			log.Printf("read disk[%s] to offset %d", disk.Disk, off)
			time.Sleep(disk.ReadDelay - time.Duration(end-begin)*time.Nanosecond)
		}
	}
	log.Printf("read disk[%s] task done", disk.Disk)
}
